#pragma once

#include <memory>

#include <glog/logging.h>

#include "base/common.h"
#include "index/primary_key_index/primary_key_file_writer/primary_key_file_writer.h"
#include "util/common_def.h"

namespace hawking {
namespace indexlib {

template <typename Key>
class SortedPrimaryKeyFileWriter : public PrimaryKeyFileWriter<Key> {
public:
    template<class K, class V>
    struct KeyValueItem {
        bool operator<(const KeyValueItem& other) {
            return key < other.key;
        }

        K key;
        V value;
    };

    using KVItem = KeyValueItem<Key, DocId>;
    SortedPrimaryKeyFileWriter() = default;
    ~SortedPrimaryKeyFileWriter() = default;

    virtual void Init(
        size_t pk_count, std::shared_ptr<NormalFileWriter> file, Pool* pool) override {
        pk_count_ = pk_count;
        pool_ = pool;
        file_ = file;
    }
    virtual Status AddPKPair(Key key, DocId docid) override;
    virtual Status AddSortedPKPair(Key key, DocId docid) override;
    virtual Status Close() override;

private:
    size_t pk_count_ = 0;
    Key last_sorted_key_;
    bool has_sorted_key_ = false;
    KVItem* buffer_ = nullptr;
    size_t pk_buffer_idx_ = 0;

    std::shared_ptr<NormalFileWriter> file_;
    Pool* pool_ = nullptr;
};

template <typename Key>
Status SortedPrimaryKeyFileWriter<Key>::AddPKPair(Key key, DocId docid) {
    if (unlikely(!buffer_)) {
        buffer_ = IE_POOL_COMPATIBLE_NEW_VECTOR(pool_, KVItem, pk_count_);
    }

    buffer_[pk_buffer_idx_].key = key;
    buffer_[pk_buffer_idx_].value = docid;
    ++pk_buffer_idx_;
    return Status::OK;
}

template <typename Key>
Status SortedPrimaryKeyFileWriter<Key>::AddSortedPKPair(Key key, DocId docid) {
    if (unlikely(!has_sorted_key_)) {
        has_sorted_key_ = true;
        last_sorted_key_ = key;
    } else {
        if (key < last_sorted_key_) {
            LOG(ERROR) << "add sort key failed. not sorted";
            return Status::InvalidArgs;
        }
    }

    KVItem item;
    item.key = key;
    item.value = docid;
    ReturnValue<size_t> ret = file_->Write((uint8_t*)&item, sizeof(item));
    return ret.Code();
}

template <typename Key>
Status SortedPrimaryKeyFileWriter<Key>::Close() {
    if (buffer_) {
        std::sort(buffer_, buffer_ + pk_count_);
        size_t buf_len = sizeof(KVItem) * pk_count_;
        file_->ReserveFile(buf_len);
        ReturnValue<size_t> ret = file_->Write((uint8_t*)buffer_, buf_len);
        if (ret.Code() != Status::OK) {
            return ret.Code();
        }

        IE_POOL_COMPATIBLE_DELETE_VECTOR(pool_, buffer_, pk_count_);
    }

    file_->Close();
    return Status::OK;
}

}
}